# Monte Carlo Experiment II: Non EPBR Data
rm(list=ls())
library(Matching)
library(MASS)
library(ebal)

# number of sims
sims <- 1000

# load lalonde data
dat <- read.table("nswre74.csv",header=T)

# unemployed
dat$u74 <- as.numeric(dat$re74==0)
dat$u75 <- as.numeric(dat$re75==0)

propensity  <- glm(treat~ I(age^2) + I(educ^2) + black +
                   hispan + married + nodegree + I(re74^2) + I(re75^2) +
                   u74 + u75, family=binomial, data=dat)

X <- as.matrix(dat[,c("age","educ","black","hispan","married","nodegree","re74","re75","u74","u75")])

M <- cbind(1,
           propensity$linear.pred,
           I(log(dat$age)^2),
           I(log(dat$educ)^2),
           I(log(dat$re74+.01)^2),
           I(log(dat$re75+.01)^2))

propensity.coeffs <- propensity$coeff

# create arbritrary weights
propensity.coeffs <- as.matrix(c(
                                 1+00,  #(Intercept)
                                 .5,    #linear.pred
                                 .01,   #age
                                 -.3,   #educ
                                 -.01, #I(re74^2)
                                 .01   #I(re75^2)
                                 ))

mu <- M %*% propensity.coeffs
Tr.pred <- exp(mu)/(1+exp(mu))
TreatmentEffect <- 1000

est    <- c("RAW","MD","GM","PS","PSMD","PSW","EB")
est.st    <- matrix(NA,sims,length(est))
colnames(est.st)   <- est
est.time  <- est.st

for(k in 1:sims){

cat("Sims no",k,"\n")

  TreatmentReal <- matrix(nrow=nrow(dat), ncol=1)
	for(i in 1:nrow(dat)){TreatmentReal[i] = sample(0:1, 1, prob=c(1-Tr.pred[i],Tr.pred[i]))}
  Y <- I(TreatmentEffect*TreatmentReal) + .1*exp(.7*log(dat$re74+0.01) + .7*log(dat$re75+0.01)) + rnorm(nrow(dat), 0, 10)
  W <- TreatmentReal

  # raw
  system.time(est.st[k,"RAW"] <- mean(Y[W==1])-mean(Y[W==0]))[3] -> est.time[k,"RAW"]
  
  # MD
  system.time(est.st[k,"MD"] <- Match(Y=Y, Tr=W, X=X, estimand="ATT",M=1,BiasAdjust = FALSE)$est )[3] ->  est.time[k,"MD"]
    
  # PS
 	system.time(pscore  <- glm(formula(paste("W~",paste(colnames(X),collapse="+"))),data=dat))[3] -> t.psc1
	system.time(est.st[k,"PS"]     <- Match(Y=Y,Tr=W, X=pscore$linear.pred,M=1,replace=FALSE,estimand="ATT",BiasAdjust = FALSE)$est)[3] -> t.psc2
  est.time[k,"PS"] <- t.psc1+t.psc2
  
  # PSMD
  L      <- cbind(1,pscore$fitted)
  Xortho <- cbind(pscore$fitted,X-(L%*%(solve(t(L)%*%L)%*%t(L)%*%X)))
  system.time(est.st[k,"PSMD"] <- Match(Y=Y,Tr=W,X=Xortho,estimand="ATT",M=1,BiasAdjust = FALSE)$est )[3] -> t.psmd
  est.time[k,"PSMD"] <- t.psc1 + t.psmd 
  
  # weighting on PS
  d <- pscore$fitted/(1-pscore$fitted)
  est.st[k,"PSW"]  <-  mean(Y[W==1]) - (sum(Y[W==0]*d[W==0])/sum(d[W==0]))
  est.time[k,"PSW"] <- t.psc1

  # GM
  suppressWarnings(system.time(gen1 <- GenMatch(X=X, Tr=W,print.level=0,estimand="ATT"))[3] ->est.time[k,"GM"])
  est.st[k,"GM"]   <- Match(Y=Y, Tr=W, X=X, Weight.matrix=gen1,estimand="ATT",M=1,BiasAdjust = FALSE)$est

  # EB
  system.time(out.eb <- ebalance(
              Treatment=W,
              X=X,
              print.level = -1
              ))[3] ->  est.time[k,"EB"]
  d <- out.eb$w
  est.st[k,"EB"]  <- mean(Y[W==1]) - (sum(Y[W==0]*d)/sum(d))

}

cat("\n\nAverage bias:\n")
colMeans(est.st)-TreatmentEffect
buf <- colMeans(est.st)-TreatmentEffect

cat("MAD\n")
apply(abs(est.st),2,median)
buf <- rbind(buf, apply(abs(est.st),2,median))

cat("RMSE:\n")
sqrt(colMeans((est.st-TreatmentEffect)^2))
buf <- rbind(buf, sqrt(colMeans((est.st-TreatmentEffect)^2)))

cat("Average time:\n")
colMeans(est.time,na.rm=TRUE)

rownames(buf) <- c("BIAS", "MAD", "RMSE")
buf <- t(buf)
buf


